#!/usr/bin/env python3

import argparse
from netCDF4 import Dataset
import numpy as np
import matplotlib.pyplot as plt

parser = argparse.ArgumentParser(description='Plot WRF terrain following vertical levels.')
parser.add_argument('-i', '--wrfout', help='WRF output file', required=True)
parser.add_argument('-o', '--figure', help='Output figure file', required=True)
args = parser.parse_args()

def demo_topo():
	nx = 101
	x = np.linspace(0, nx - 1, nx) - (nx - 1) / 2
	zs = 5000 * np.exp(-0.01 * x**2)
	p0 = 1000e2
	t0 = 288
	return (p0, t0, x, zs)

def plot_vert_levels(wrfout_file, figure_file):
	f    = Dataset(wrfout_file, 'r')
	etac = f.ETAC
	p00  = f.variables['P00'][0]
	ptop = f.variables['P_TOP'][0]
	znu  = f.variables['ZNW'][0]
	c1   = 2 * etac**2 / (1 - etac)**3
	c2   = -etac * (4 + etac + etac**2) / (1 - etac)**3
	c3   = 2 * (1 + etac + etac**2) / (1 - etac)**3
	c4   = -(1 + etac) / (1 - etac)**3
	b    = c1 + c2 * znu + c3 * znu**2 + c4 * znu**3
	for k in range(len(znu)):
		if znu[k] < etac: b[k] = 0

	p0, t0, x, zs = demo_topo()

	Rd = 287
	ps = p0 * np.exp(-zs / Rd / t0)
	p  = np.ndarray((len(znu), len(x)))
	for k in range(len(znu)):
		p[k,:]  = b[k] * (ps - ptop) + (znu[k] - b[k]) * (p00 - ptop) + ptop

	ax = plt.subplot(111)
	for k in range(len(znu)):
		ax.plot(x, p[k,:] / 100)
	ax.fill_between(x, ps / 100, p00 / 100, color='black', facecolor='black')
	ax.set_ylabel('Pressure (hPa)')
	ax.invert_yaxis()
	plt.savefig(figure_file)
	plt.close()

plot_vert_levels(args.wrfout, args.figure)

